ENH: implement `at` #53

crusaderky · 2024-12-09T13:49:26Z

Implement a new at(x, idx) or at(x)[idx] function, mocking the syntax of JAX's omonymous method .
This is propaedeutic to JAX support in libraries that support the Array API, e.g. scipy.

Moved from data-apis/array-api-compat#205

Blockers

Blocks

ENH: _lib: JAX support (non-jitted) scipy/scipy#22070

pyproject.toml

tests/test_at.py

pyproject.toml

src/array_api_extra/_funcs.py

crusaderky · 2024-12-10T12:44:38Z

Rebased on top of #58

crusaderky · 2024-12-11T16:09:32Z

Reverted rebase onto Saner linters #58
Fixed all linters issues
Added unit tests for all backends
Removed **kwargs. Their main original appeal was allowing to pass performance tweaks to JAX, which are indices_are_sorted=True and unique_indices=True. However, I have now realised that the array API specification does not support integer array indices, at all, and both of these parameters are pointless without.
I'm having big second thoughts on get() - read below.

crusaderky · 2024-12-11T16:15:35Z

Any idea how to fix this? Looks like the ubuntu_latest VM has an obsolete driver (or more likely no driver)

ERROR cupy_backends.cuda.api.runtime.CUDARuntimeError: cudaErrorInsufficientDriver: CUDA driver version is insufficient for CUDA runtime version

src/array_api_extra/_funcs.py

lucascolley · 2024-12-11T16:57:43Z

I don't think we can get GPU CI without paying someone for it, cc @rgommers .

crusaderky · 2024-12-11T17:22:56Z

I don't think we can get GPU CI without paying someone for it, cc @rgommers .

Can cupy run on a CPU-only host?
I've disabled it in CI for the time being

rgommers · 2024-12-11T19:48:18Z

I don't think we can get GPU CI without paying someone for it, cc @rgommers .

Yep, that's on my radar to push forward this month, on multiple projects. Please feel free to open a new issue and assign it to me. I think we can hook up a shared GPU runner between this project and array-api-compat.

Can cupy run on a CPU-only host?

I don't think so.

pyproject.toml

crusaderky · 2024-12-30T13:06:55Z

@lucascolley ready for merge and release! 🥳

lucascolley · 2024-12-30T13:15:49Z

good to proceed @rgommers ?

src/array_api_extra/_funcs.py

lucascolley · 2025-01-01T23:22:01Z

looks great, thanks! Note to self: we ought to add

# sphinx-copybutton configurations
copybutton_prompt_text = r">>> |\.\.\. |\$ |In \[\d*\]: | {2,5}\.{3,}: | {5,8}: "
copybutton_prompt_is_regexp = True

to docs/conf.py.

crusaderky · 2025-01-02T10:21:57Z

Annoyingly, codecov is glitching (again). It's stating in the "Files" view of this PR that the tests never run under JAX, but if you follow the link from the log of Check ci-backends it says that everything's green.

lucascolley · 2025-01-02T10:30:00Z

yeah, happy to just ignore that, but also happy if anyone figures out how to fix that

rgommers

Thanks @crusaderky and reviewers. This is looking quite good. A couple of comments (disclaimer: it's a little hard to dig through the already-resolved review comments, so I may be saying something that was already discussed there):

(1) The copy=True default seems wrong to me. The main point of this feature is that we can replace in-place operations like x[idx, :] += y in SciPy et al. with something new to support JAX. That something new should be the idiomatic JAX way, so x = xpx.at(x)[idx, :].add(y). This should come with as little of a performance penalty for NumPy usage as possible - hence the copy=True default doesn't work. It's already a significant concession to code readability to have to avoid += & co, and having to also add copy=None in all places where performance matters seems too much.

(2) When playing around with this branch, most things work as expected, but this was a little surprising:

>>> x = np.arange(5)
>>> x = xpx.at(x)[[3, 4]].divide(1.5)
>>> x.dtype
dtype('int64')
>>> x = np.arange(5)
>>> x /= 1.5
Traceback (most recent call last):
  Cell In[32], line 1
    x /= 1.5
UFuncTypeError: Cannot cast ufunc 'divide' output from dtype('float64') to dtype('int64') with casting rule 'same_kind'

Is this on purpose or by accident? Raising on these cases may be cleaner if that can be done in a performant way.

(3) Library support is pretty comprehensive already; a few others come to mind. MLX is already supported by the generic path since it allows in-place syntax (I think, didn't verify). It also has array.at which matches JAX syntax/semantics. ndonnx is going to need explicit support in the future as well, since it's immutable like JAX, but that can be left for later since it doesn't seem to have Array.at or another alternative here. tl;dr nothing to change in this PR I think.

rgommers · 2025-01-02T20:49:01Z

src/array_api_extra/_funcs.py

+        Allow for the alternate syntax ``at(x)[start:stop:step]``.
+
+        It looks prettier than ``at(x, slice(start, stop, step))``
+        and feels more intuitive coming from the JAX documentation.


This should probably be the default way of doing things in the examples? Best to steer people to what feels idiomatic to JAX users I'd think. I'd de-emphasize, or perhaps even leave out completely (if there are no downsides to that), the alternative syntax.

I've been writing several prototype PRs on various submodules of Sphinx.
From the experience I gathered there, I must say that at(x)[idx] feels more natural and easier on the eyes when there are multiple indices and/or slices, whereas at(x, idx) feels nicer and less cluttered in all other cases.

Okay, fair enough - seems like an argument to keep both then.

In fact, I've been considered writing aliases for the benefit of readability set_at(x, idx, value) etc. but haven't gone for them for the sake of keeping the API surface small.

Not adding a third way sounds like a healthy idea at least for now - let's see how we like this in its current form I'd say.

crusaderky · 2025-01-03T12:34:00Z

(1) The copy=True default seems wrong to me.

Changed it to copy=None

(2) x /= 1.5

TIL that

x[idx] = x[idx].__add__(value)
x[idx] = x[idx].__iadd__(value)

are almost always equivalent!

Fixed.

tests/test_at.py

lucascolley · 2025-01-03T12:51:40Z

shall we get going with a release? :)

rgommers

Looks ready to me now, +1 for merge & release if everyone else is happy too.

lucascolley · 2025-01-03T13:24:46Z

let's give this a whirl, thanks all! if you can point the SciPy PR to the release when it is out @crusaderky, I can update the sklearn PR

crusaderky commented Dec 9, 2024

View reviewed changes

pyproject.toml Outdated Show resolved Hide resolved

crusaderky commented Dec 9, 2024

View reviewed changes

tests/test_at.py Outdated Show resolved Hide resolved

lucascolley reviewed Dec 9, 2024

View reviewed changes

tests/test_at.py Outdated Show resolved Hide resolved

lucascolley reviewed Dec 9, 2024

View reviewed changes

pyproject.toml Outdated Show resolved Hide resolved

lucascolley marked this pull request as ready for review December 9, 2024 14:30

lucascolley marked this pull request as draft December 9, 2024 14:31

lucascolley reviewed Dec 9, 2024

View reviewed changes

src/array_api_extra/_funcs.py Outdated Show resolved Hide resolved

crusaderky force-pushed the at branch from 162264b to dd068d2 Compare December 9, 2024 17:13

lucascolley added enhancement New feature or request new function labels Dec 9, 2024

crusaderky force-pushed the at branch 9 times, most recently from 376ad49 to 4a0943a Compare December 10, 2024 12:36

crusaderky mentioned this pull request Dec 10, 2024

Saner linters #58

Closed

lucascolley mentioned this pull request Dec 10, 2024

POC: appease linter for gh-53 #59

Closed

crusaderky force-pushed the at branch from 4a0943a to 2b18097 Compare December 11, 2024 16:05

crusaderky commented Dec 11, 2024

View reviewed changes

src/array_api_extra/_funcs.py Outdated Show resolved Hide resolved

crusaderky force-pushed the at branch from a0e20a9 to 3c0a43b Compare December 11, 2024 17:22

lucascolley mentioned this pull request Dec 26, 2024

ENH: pad: add delegation #72

Merged

8 tasks

ENH: add new function at

34f89f5

lucascolley force-pushed the at branch from 825b68c to 34f89f5 Compare December 26, 2024 19:20

lucascolley added 2 commits December 26, 2024 19:26

MAINT: use released array-api-compat

a50c5ae

update lock-file

f19e42f

lucascolley force-pushed the at branch from 68ae073 to f19e42f Compare December 26, 2024 20:09

lucascolley approved these changes Dec 26, 2024

View reviewed changes

Update dependencies

64396d0

crusaderky commented Dec 30, 2024

View reviewed changes

pyproject.toml Show resolved Hide resolved

crusaderky commented Dec 30, 2024

View reviewed changes

pyproject.toml Show resolved Hide resolved

lucascolley reviewed Jan 1, 2025

View reviewed changes

src/array_api_extra/_funcs.py Outdated Show resolved Hide resolved

Add xpx namespace in documentation

98141ca

rgommers self-requested a review January 2, 2025 20:38

rgommers reviewed Jan 2, 2025

View reviewed changes

Change copy to default to None

38562db

raise on incompatible cast

b19cc11

crusaderky force-pushed the at branch from 2edc44a to b19cc11 Compare January 3, 2025 12:36

lucascolley reviewed Jan 3, 2025

View reviewed changes

tests/test_at.py Outdated Show resolved Hide resolved

Update tests/test_at.py

26f9a7d

rgommers approved these changes Jan 3, 2025

View reviewed changes

lucascolley merged commit 84bf725 into data-apis:main Jan 3, 2025
10 checks passed

crusaderky deleted the at branch January 3, 2025 13:26

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

ENH: implement `at` #53

ENH: implement `at` #53

crusaderky commented Dec 9, 2024 •

edited

Loading

crusaderky commented Dec 10, 2024

crusaderky commented Dec 11, 2024 •

edited

Loading

crusaderky commented Dec 11, 2024 •

edited

Loading

lucascolley commented Dec 11, 2024

crusaderky commented Dec 11, 2024

rgommers commented Dec 11, 2024

crusaderky commented Dec 30, 2024

lucascolley commented Dec 30, 2024

lucascolley commented Jan 1, 2025

crusaderky commented Jan 2, 2025 •

edited

Loading

lucascolley commented Jan 2, 2025

rgommers left a comment

rgommers Jan 2, 2025

crusaderky Jan 3, 2025 •

edited

Loading

rgommers Jan 3, 2025

crusaderky Jan 3, 2025

rgommers Jan 3, 2025

crusaderky commented Jan 3, 2025

lucascolley commented Jan 3, 2025

rgommers left a comment

lucascolley commented Jan 3, 2025

ENH: implement at #53

ENH: implement at #53

Conversation

crusaderky commented Dec 9, 2024 • edited Loading

Blockers

Blocks

crusaderky commented Dec 10, 2024

crusaderky commented Dec 11, 2024 • edited Loading

crusaderky commented Dec 11, 2024 • edited Loading

lucascolley commented Dec 11, 2024

crusaderky commented Dec 11, 2024

rgommers commented Dec 11, 2024

crusaderky commented Dec 30, 2024

lucascolley commented Dec 30, 2024

lucascolley commented Jan 1, 2025

crusaderky commented Jan 2, 2025 • edited Loading

lucascolley commented Jan 2, 2025

rgommers left a comment

Choose a reason for hiding this comment

rgommers Jan 2, 2025

Choose a reason for hiding this comment

crusaderky Jan 3, 2025 • edited Loading

Choose a reason for hiding this comment

rgommers Jan 3, 2025

Choose a reason for hiding this comment

crusaderky Jan 3, 2025

Choose a reason for hiding this comment

rgommers Jan 3, 2025

Choose a reason for hiding this comment

crusaderky commented Jan 3, 2025

lucascolley commented Jan 3, 2025

rgommers left a comment

Choose a reason for hiding this comment

lucascolley commented Jan 3, 2025

ENH: implement `at` #53

ENH: implement `at` #53

crusaderky commented Dec 9, 2024 •

edited

Loading

crusaderky commented Dec 11, 2024 •

edited

Loading

crusaderky commented Dec 11, 2024 •

edited

Loading

crusaderky commented Jan 2, 2025 •

edited

Loading

crusaderky Jan 3, 2025 •

edited

Loading